import argparse
import logging
import os
import random
import sys
import time
from datetime import datetime

import torch


def get_config():
    """
    配置训练参数

    :returns args : 解析的参数
    :returns logger : 日志记录器
    """
    # 创建命令行参数解析器
    parser = argparse.ArgumentParser()

    """Base"""
    # 添加模型类别参数，默认是2个类别
    parser.add_argument('--num_classes', type=int, default=2)

    # 添加基座模型参数，默认为'bert'，可选'distilbert'
    parser.add_argument('--base_model', type=str, default='bert',
                        choices=['bert','distilbert'])

    # 添加方法，默认为'fnn'，可选多种方法如'gru', 'rnn', 'lstm'等
    parser.add_argument('--method', type=str, default='fnn',
                        choices=['fnn', 'gru', 'rnn', 'bilstm', 'lstm'])

    """Optimization"""
    # 添加训练批量大小参数，默认值为16
    parser.add_argument('--train_batch_size', type=int, default=16)

    # 添加验证批量大小参数，默认值为16
    parser.add_argument('--eval_batch_size', type=int, default=16)

    # 添加训练轮次(epoch)参数，默认值为2
    parser.add_argument('--num_epoch', type=int, default=2)

    # 添加学习率参数，默认值为5e-5
    parser.add_argument('--lr', type=float, default=5e-5)

    # 添加dropout参数，默认值为0.1
    parser.add_argument('--dropout', type=float, default=0.1)

    '''Environment'''
    # 添加设备参数，默认为'cpu'，可以选择使用'cuda'（GPU）或者'cpu'
    parser.add_argument('--device', type=str, default='cuda',
                        choices=['cpu', 'cuda'])

    # 添加一个时间戳参数，用于标识日志文件等，默认生成一个由时间和随机数组成的字符串
    parser.add_argument('--timestamp', type=int,
                        default=f'{time.time():.0f}{random.randint(0, 999):03}')

    # 解析传递的命令行参数
    args = parser.parse_args()

    """Logger"""
    # 根据模型名称、方法名称和当前时间生成日志文件名
    args.log_name = f'{args.base_model}_{args.method}_{datetime.now().strftime("%Y_%m_%d_%H_%M_%S")[2:]}.log'

    # 如果logs目录不存在，则创建该目录
    if not os.path.exists('logs'):
        os.mkdir('logs')

    # 设置日志记录器
    logger = logging.getLogger()
    logger.setLevel(logging.INFO)

    # 添加日志处理程序，将日志信息输出到标准输出（控制台）
    logger.addHandler(logging.StreamHandler(sys.stdout))

    # 添加日志处理程序，将日志信息保存到文件
    logger.addHandler(logging.FileHandler(os.path.join('logs', args.log_name)))

    # 返回解析的参数和日志记录器
    return args, logger


# 测试 get_config 函数
if __name__ == "__main__":
    # 获取配置和日志记录器
    args, logger = get_config()

    # 打印参数
    print("Parsed Arguments:")
    for arg in vars(args):
        print(f"{arg}: {getattr(args, arg)}")

    # 使用日志记录器打印一条测试消息
    logger.info("\rThis is a test log message.\rThis is a test log message.")
